#!/usr/bin/python2.4

"""KD-tree for nearest neighbor  search in a k-dimensional space.

See:
  http://code.activestate.com/recipes/
    577497-kd-tree-for-nearest-neighbor-search-in-a-k-dimensi/
"""

import collections
import itertools
import math
import matrix
import random

import namedtuple

# Node = collections.namedtuple('Node', 'point axis label left right')
Node = namedtuple.namedtuple('Node', 'point axis label left right')

class KDTree(object):
  """A tree for nearest neighbor search in a k-dimensional space.

  For information about the implementation, see
    http://en.wikipedia.org/wiki/Kd-tree

  Usage:
  objects is an iterable of (point, label) tuples
  k is the number of dimensions

  t = KDTree(k, objects)
  point, label, distance = t.NearestNeighbor(destination)
  """

  def __init__(self, k, objects=[]):

    def _BuildTree(objects, axis=0):
      if not objects:
        return None

      objects.sort(key=lambda o: o[0].KeyByAxis(axis))
      median_idx = len(objects) // 2
      median_point, median_label = objects[median_idx]

      next_axis = (axis + 1) % k
      return Node(median_point,
                  axis,
                  median_label,
                  _BuildTree(objects[:median_idx], next_axis),
                  _BuildTree(objects[median_idx + 1:], next_axis))

    self.root = _BuildTree(list(objects))

  def NearestNeighbor(self, destination):
    """Compute the nearest neighbor to the destination.

    Args:
      destination: Our index (?).

    Returns:
      A triple: p1, p2 and the distance between them.
    """

    # state of search
    best = [
        None,  # best point found
        None,  # its label
        float('inf')  # lowest squared distance
        ]

    def _RecursiveSearch(here):
      """Search for the point starting from the node 'here'."""
      if here is None:
        return
      point, axis, label, left, right = here

      here_sd = point.SquareDistance(destination)
      if here_sd < best[2]:
        best[:] = point, label, here_sd

      # diff = destination[axis] - point[axis]
      diff = destination.KeyByAxis(axis) - point.KeyByAxis(axis)
      # close, away = (left, right) if diff <= 0 else (right, left)
      if diff <= 0:
        close, away = (left, right)
      else:
        close, away = (right, left)

      _RecursiveSearch(close)
      if diff ** 2 < best[2]:
        _RecursiveSearch(away)

    _RecursiveSearch(self.root)
    return best[0], best[1], math.sqrt(best[2])


def main():
  random.seed(3141592653)
  k = 2
  npoints = 10
  lookups = 10
  eps = 1e-8

  points = [(matrix.LatLon(random.random(), random.random()), i)
            for i in xrange(npoints)]

  tree = KDTree(k, points)

  for _ in xrange(lookups):
    destination = matrix.LatLon(random.random(), random.random())
    point, label, mindistance = tree.NearestNeighbor(destination)

    minsq = min(p.SquareDistance(destination) for p, _ in points)
    if True:
      print 'point=%s label=%s sqrt(minsq)=%6.4f mindistance=%6.4f' % (
          point, label, math.sqrt(minsq), mindistance)
    assert abs(math.sqrt(minsq) - mindistance) < eps


if __name__ == '__main__':
  main()
